'''
Total reward:  9862.0
Mean reward:  98.62
'''
import gymnasium as gym
import numpy as np
import ale_py
from lib import common


gym.register_envs(ale_py)

# Initialize the environment
# env = common.ProcessFrame84(gym.make('ALE/Atlantis2-v5', render_mode="human", frameskip=1))
env = gym.make('ALE/Atlantis2-v5')

# Reset the environment to get the initial state
total_reward = 0
# Run a loop to play the game
episo = 100
for _ in range(episo):
    state = env.reset()
    for _ in range(10000):
        # env.render()  # Render the environment

        # Take a random action
        action = env.action_space.sample()

        # Get the next state, reward, done flag, and info from the environment
        state, reward, done, trunc, info = env.step(action)
        total_reward += reward
        if reward != 0:
            print("reward: ", reward)
            print("info: ", info)

        # If done, reset the environment
        if done or trunc:
            break

print("Total reward: ", total_reward)
print("Mean reward: ", total_reward / episo)

# Close the environment
env.close()